#pragma once
#include "util.hpp"
#include <atomic>
#include <cassert>
#include <deque>
#include <functional>
#include <mutex>
#include <optional>
#include <span>
#include <thread>

namespace Thread::util {

// 全局队列 基于互斥锁，无条件变量
class GlobalQueue : Nocopyable {
public:
  explicit GlobalQueue() = default;

public:
  [[nodiscard]]
  bool closed() {
    return __closed.load(std::memory_order::acquire);
  }
  void close() {
    if (closed())
      return;
    __closed.store(true, std::memory_order::release);
  }
  [[nodiscard]]
  std::size_t size() {
    std::lock_guard lock{__mtx};
    return __queue.size();
  }

  [[nodiscard]]
  bool empty() {
    std::lock_guard lock{__mtx};
    return __queue.empty();
  }

public:
  void push(std::function<void()> f) {
    std::lock_guard lock{__mtx};
    if (closed())
      return;
    __queue.push_back(std::move(f));
  }

  void push_batch(std::vector<std::function<void()>> &&funcs) {
    std::lock_guard lock{__mtx};
    for (auto &f : funcs) {
      if (closed())
        return;
      __queue.push_back(f);
    }
  }
  std::optional<std::function<void()>> try_pop() {
    std::lock_guard lock{__mtx};
    if (__queue.empty()) {
      return std::nullopt;
    }
    auto result = std::move(__queue.front());
    __queue.pop_front();
    return result;
  }

  std::vector<std::function<void()>> try_pop_batch(std::size_t n) {
    std::lock_guard lock{__mtx};
    n = std::min(__queue.size(), n);
    std::vector<std::function<void()>> result;
    for (size_t i = 0; i < n; ++i) {
      result.push_back(std::move(__queue.front()));
      __queue.pop_front();
    }
    return result;
  }

private:
  std::mutex __mtx{};
  std::deque<std::function<void()>> __queue;
  std::atomic<bool> __closed{false};
};

// 本地队列
// 基于无锁环形队列，支持任务窃取，支持队满负载均衡。总体是一个单生产者多消费者（单写多读）的场景。
template <std::size_t CAPACITY = 256> class LocalQueue : Nocopyable {
  static_assert((CAPACITY & (CAPACITY - 1)) == 0,
                "CAPACITY must be power of 2");
  static_assert(CAPACITY > 0, "CAPACITY must be greater than 0");

public:
  LocalQueue() = default;

  ~LocalQueue() {}

public:
  [[nodiscard]]
  std::size_t remain_size() {
    auto [steal, _] = unpack(__consumer.load(std::memory_order::acquire));
    auto write = __producer.load(std::memory_order::acquire);
    assert(CAPACITY >= static_cast<std::size_t>(write - steal));
    return CAPACITY - static_cast<std::size_t>(write - steal);
  }

  [[nodiscard]]
  std::uint32_t size() {
    auto [_, read] = unpack(__consumer.load(std::memory_order::acquire));
    auto write = __producer.load(std::memory_order::acquire);
    return write - read;
  }

  [[nodiscard]]
  bool empty() {
    return size() == 0u;
  }

  [[nodiscard]]
  std::size_t capacity() const {
    return __data.max_size();
  }

public:
  // 单个元素push_back（从尾部push） ,不处理队满：直接返回false
  bool push_back(std::function<void()> func) {
    // 解包获取到偷取指针
    auto [steal, _] = unpack(__consumer.load(std::memory_order::acquire));
    // 获取到写指针
    auto write = __producer.load(std::memory_order::relaxed);
    // 如果满了，返回false
    if (write - steal > static_cast<uint32_t>(capacity())) {
      return false;
    }
    // 先写数据
    __data[write & mask] = std::move(func);
    // 再更新写指针
    write += 1;
    __producer.store(write, std::memory_order::release);
    return true;
  }
  // 批量push-不处理队满直接抛异常
  void push_back_batch(std::span<std::function<void()>> funcs) {
    for (auto &f : funcs) {
      if (!push_back(f)) {
        throw std::runtime_error("push_back overflow!");
      }
    }
  }

  // 考虑队满的单个push
  void push_back_maybe_overflow(std::function<void()> func,
                                GlobalQueue &glbqueue) {
    uint32_t write{0};
    while (true) {
      auto consumer = __consumer.load(std::memory_order::acquire);
      auto [steal, read] = unpack(consumer);
      write = __producer.load(std::memory_order::relaxed);
      if (write - steal < static_cast<uint32_t>(CAPACITY)) {
        // 未满
        break;
      } else if (steal != read) {
        // 有其他线程再窃取,直接放入全局队列
        glbqueue.push(func);
        return;
      } else {
        // 正常处理队满
        if (push_already_overflow(func, read, write, glbqueue)) {
          return;
        }
      }
    }
    std::size_t idx = static_cast<std::size_t>(write) & mask;
    __data[idx] = std::move(func);
    __producer.store(write + 1, std::memory_order::release);
  }

  std::optional<std::function<void()>> pop() {
    auto cur_consumer = __consumer.load(std::memory_order::acquire);
    std::size_t index = 0;
    while (true) {
      auto [cur_steal, cur_read] = unpack(cur_consumer);
      auto cur_write = __producer.load(std::memory_order::relaxed);
      // 队空
      if (cur_read == cur_write) {
        return std::nullopt;
      }
      // 构造下一次的指针
      auto next_read = cur_read + 1;
      auto next_consumer = cur_steal == cur_read ? pack(next_read, next_read)
                                                 : pack(cur_steal, next_read);
      if (__consumer.compare_exchange_weak(cur_consumer, next_consumer,
                                           std::memory_order::acq_rel,
                                           std::memory_order::acquire)) {
        index = static_cast<std::size_t>(cur_read) & mask;
        break;
      }
    }
    return __data[index];
  }

  // 窃取接口。本地队列窃取目标队列
  /*
   *   窃取规则：
   *   1.如果有其他线程在窃取就直接返回
   *   2.如果目标队列任务数大于一半，窃取目标队列的一半，否则直接返回
   */
  std::optional<std::function<void()>> steal_to(LocalQueue &dst) {
    std::optional<std::function<void()>> result{std::nullopt};

    // 获取目标队列的窃取指针和写指针
    auto [dst_steal, _] =
        unpack(dst.__consumer.load(std::memory_order::acquire));
    auto dst_write = dst.__producer.load(std::memory_order::relaxed);

    // 如果没有一半，直接返回
    if (dst_write - dst_steal > static_cast<uint32_t>(CAPACITY / 2)) {
      return result;
    }

    // 进行窃取
    auto n = steal_into_impl(dst, dst_write);
    // 如果没有任务直接返回
    if (n == 0) {
      return result;
    }
    // 将最后一个任务做完结构返回
    n -= 1;
    // 构造新的写指针
    auto dst_new_write = dst_write + n;
    std::size_t index = static_cast<std::size_t>(dst_new_write) & mask;
    // 拿出最后一个任务放入返回值
    result = std::move(dst.__data[index]);
    if (n > 0) {
      // 更新写指针
      dst.__producer.store(dst_new_write, std::memory_order::release);
    }
    return result;
  }

private:
  // 队满后单个元素push,队满后处理逻辑:将本地队列一半的任务转移到全局队列。
  /*注意:
   *    1.调用的时候队已经满了
   *    2.调用的时候确保没有其他线程进行窃取。
   */

  bool push_already_overflow(std::function<void()> func, std::uint32_t rd,
                             std::uint32_t wt, GlobalQueue &glbqueue) {
    // 获取到队列一半的大小作为转移的数量
    auto take_len = static_cast<uint32_t>(CAPACITY / 2);
    assert(wt - rd == CAPACITY);

    // 打包当前整个消费者指针
    auto cur_consumer = pack(rd, rd);
    // 打包下一个整个消费者指针
    auto next_consumer = pack(rd + take_len, rd + take_len);
    // 更新读指针，如果cas失败返回false
    if (!__consumer.compare_exchange_weak(cur_consumer, next_consumer,
                                          std::memory_order::release,
                                          std::memory_order::relaxed)) {
      return false;
    }

    // 构造一个vector，从队列里把数据放入vector
    std::vector<std::function<void()>> funcs;
    for (int i = 0; i < take_len; ++i) {
      std::size_t idx = static_cast<std::size_t>(rd + i) & mask;
      funcs.push_back(std::move(__data[idx]));
    }
    // 把当前任务也放入vector里
    funcs.push_back(std::move(func));
    // 向全局队列批量push
    glbqueue.push_batch(std::move(funcs));
    return true;
  }
  // 主窃取逻辑
  /*
   *   窃取规则：
   *   1.如果有其他线程在窃取就直接返回
   *   2.如果目标队列任务数大于一半，窃取目标队列的一半，否则直接返回
   */

  std::uint32_t steal_into_impl(LocalQueue &dst, uint32_t dst_wt) {
    // 获取当前完整消费者指针
    std::uint64_t cur_consumer = __consumer.load(std::memory_order::acquire);
    // 声明下一个完整指针变量
    std::uint64_t next_consumer;
    uint32_t steal_num = 0;
    while (true) {
      // 获取到当前队列（源队列）的窃取指针和读指针
      auto [src_steal, src_read] = unpack(cur_consumer);
      // 获取到当前队列（源队列）的写指针
      auto src_write = __producer.load(std::memory_order::acquire);
      // 窃取指针如果和读指针不一样，说明其他线程正在窃取，那直接退出
      if (src_steal != src_read) {
        return 0;
      }
      // 获取到队列现有任务数的一半
      steal_num = (src_write - src_read) / 2;
      if (steal_num == 0) {
        return 0;
      }

      // 构造并更新下一次的消费者指针，告诉其他线程，我开始窃取了（窃取指针不等于读指针）
      auto next_src_read = src_read + steal_num;
      assert(src_steal != next_src_read);
      next_consumer = pack(src_steal, next_src_read);
      if (__consumer.compare_exchange_weak(cur_consumer, next_consumer,
                                           std::memory_order::acq_rel,
                                           std::memory_order::acquire)) {
        break;
      }
    }
    // 获取到当前的窃取指针
    auto [first, _] = unpack(next_consumer);
    // 将目标队列的任务转移到当前队列
    for (uint32_t i = 0; i < steal_num; ++i) {
      auto src_idx = static_cast<std::size_t>(first + i) & mask;
      auto dst_idx = static_cast<std::size_t>(dst_wt + i) & mask;
      dst.__data[dst_idx] = std::move(__data[src_idx]);
    }
    // 重新获取到当前指针（因为之前的cur指针还没有更新成现在的指针）
    cur_consumer = next_consumer;
    // 构造并更新消费者指针，恢复无线程窃取状态（窃取指针等于读指针）
    while (true) {
      // 获取到当前指针的读指针
      auto [_, rd] = unpack(cur_consumer);
      // 构造下一次消费者指针，让窃取指针和读指针一致恢复无线程窃取状态
      next_consumer = pack(rd, rd);
      if (__consumer.compare_exchange_weak(cur_consumer, next_consumer,
                                           std::memory_order::acq_rel,
                                           std::memory_order::acquire)) {
        return steal_num;
      }
    }
  }

private:
  static constexpr std::size_t mask = CAPACITY - 1;

  /*
   * 功能:将64位消费者指针拆成两部分，每个部分32位
   * 高32位代表窃取指针，低32位代表读指针
   * 实现原理 ：
   *   packed >> 32 ：右移32位，获取高32位的值
   *   static_cast<uint32_t>(packed) ：直接转换为32位整数，自动截取低32位
   */
  [[nodiscard]]
  static auto unpack(uint64_t packed) -> std::pair<uint32_t, uint32_t> {
    return {static_cast<uint32_t>(packed >> 32), static_cast<uint32_t>(packed)};
  }

  /*
   * 功能 ：将两个32位整数合并成一个64位整数
   * 实现原理 ：
   *    static_cast<uint64_t>(steal) << 32
   * ：将第一个参数转换为64位并左移32位，放在高32位 static_cast<uint64_t>(rd)
   * ：将第二个参数转换为64位，占据低32位 使用按位或操作符 | 将两部分合并
   */
  [[nodiscard]]
  static auto pack(uint32_t steal, uint32_t rd) -> uint64_t {
    return static_cast<uint64_t>(steal) << 32 | static_cast<uint64_t>(rd);
  }

private:
  std::atomic<std::uint64_t> __consumer; // 消费者指针，拆成两部分防止ABA
  std::atomic<std::uint32_t>
      __producer; // 生产者指针，由于是单生产者场景，可以不用防范ABA
  std::array<std::function<void()>, CAPACITY> __data; // 数据
};
} // namespace Thread::util
